package main.java.com.pzhu.test;

import java.util.concurrent.atomic.AtomicLong;

public class SnowflakeIdGenerator {

    // 起始时间戳（毫秒）
    private static final long START_TIMESTAMP = 1609459200000L;

    // 工作机器 ID（可根据实际情况设置）
    private final long workerId;

    // 数据中心 ID（可根据实际情况设置）
    private final long dataCenterId;

    // 每一部分的位数
    private static final int WORKER_ID_BITS = 10;
    private static final int DATA_CENTER_ID_BITS = 10;
    private static final int SEQUENCE_BITS = 12;

    // 最大工作机器 ID
    private static final long MAX_WORKER_ID = -1L ^ (-1L << WORKER_ID_BITS);
    // 最大数据中心 ID
    private static final long MAX_DATA_CENTER_ID = -1L ^ (-1L << DATA_CENTER_ID_BITS);

    // 工作机器 ID 左移位数
    private static final long WORKER_ID_SHIFT = SEQUENCE_BITS;
    // 数据中心 ID 左移位数
    private static final long DATA_CENTER_ID_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS;
    // 时间戳左移位数
    private static final long TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS + DATA_CENTER_ID_BITS;

    // 上次生成 ID 的时间戳
    private final AtomicLong lastTimestamp = new AtomicLong(-1L);

    // 序列号
    private final AtomicLong sequence = new AtomicLong(0L);

    public SnowflakeIdGenerator(long workerId, long dataCenterId) {
        if (workerId < 0 || workerId > MAX_WORKER_ID) {
            throw new IllegalArgumentException("工作机器 ID 超出范围");
        }
        if (dataCenterId < 0 || dataCenterId > MAX_DATA_CENTER_ID) {
            throw new IllegalArgumentException("数据中心 ID 超出范围");
        }
        this.workerId = workerId;
        this.dataCenterId = dataCenterId;
    }

    public synchronized long generateId() {
        long currentTimestamp = System.currentTimeMillis();

        // 时钟回拨
        if (currentTimestamp < lastTimestamp.get()) {
            throw new RuntimeException("时钟回拨");
        }

        // 相同时间戳，序列号递增
        if (currentTimestamp == lastTimestamp.get()) {
            sequence.compareAndSet((1 << SEQUENCE_BITS) - 1, 0);
            sequence.incrementAndGet();
        } else {
            sequence.set(0L);
        }

        lastTimestamp.set(currentTimestamp);

        return ((currentTimestamp - START_TIMESTAMP) << TIMESTAMP_SHIFT)
                | (dataCenterId << DATA_CENTER_ID_SHIFT)
                | (workerId << WORKER_ID_SHIFT)
                | sequence.get();
    }

    public static void main(String[] args) {
        SnowflakeIdGenerator generator = new SnowflakeIdGenerator(1, 1);
        for (int i = 0; i < 10; i++) {
            System.out.println(generator.generateId());
        }
    }
}